from easyocr import Reader
from pdf2image import convert_from_path
import os
import json
from datetime import datetime


# TODO arm架构适配
class PDFOCRProcessor:
    def __init__(self, use_gpu=True, debug=True):
        """
        初始化OCR处理器
        :param use_gpu: 是否使用GPU
        :param debug: 是否打印调试信息
        """
        self.reader = Reader(['ch_sim', 'en'], gpu=use_gpu)
        self.debug = debug

    def convert_pdf_to_images(self, pdf_path, output_dir=None):
        if output_dir is None:
            output_dir = f"easyocr_temp_images_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        os.makedirs(output_dir, exist_ok=True)

        images = convert_from_path(pdf_path, dpi=600)
        image_paths = []
        for i, image in enumerate(images):
            image_path = os.path.join(output_dir, f'page_{i + 1}.jpg')
            image.save(image_path, 'JPEG')
            image_paths.append(image_path)

        return image_paths, output_dir

    def process_image(self, image_path):
        if self.debug:
            print(f"\n处理图像: {image_path}")

        result = self.reader.readtext(image_path, detail=1)
        if self.debug:
            for bbox, text, confidence in result:
                print(f"识别文字: {text} (置信度: {confidence})")
                print(f"边界框: {bbox}")

        processed_result = []
        for bbox, text, confidence in result:
            if isinstance(bbox, list):
                bbox = [float(coord) for point in bbox for coord in point]
            else:
                raise ValueError(f"Unexpected bbox format: {bbox}")
            confidence = float(confidence)  # 转换为 Python 原生 float
            processed_result.append({
                "type": "text",
                "content": text,
                "bbox": bbox,
                "confidence": confidence
            })

        return processed_result

    def save_results(self, results, output_path):
        processed_results = []
        for page_num, page_result in enumerate(results, 1):
            page_data = {"page_number": page_num, "regions": page_result}
            processed_results.append(page_data)

        with open(output_path + '.json', 'w', encoding='utf-8') as f:
            json.dump(processed_results, f, ensure_ascii=False, indent=2)

        # with open(output_path + '.txt', 'w', encoding='utf-8') as f:
        #     for page in processed_results:
        #         f.write(f"\n=== 第 {page['page_number']} 页 ===\n\n")
        #         for region in page["regions"]:
        #             if region["type"] == "text":
        #                 f.write(f"{region['content']}\n")

    def process_pdf(self, pdf_path, output_path):
        try:
            if self.debug:
                print(f"\n开始处理PDF: {pdf_path}")

            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            image_paths, temp_dir = self.convert_pdf_to_images(pdf_path)
            if self.debug:
                print(f"转换得到 {len(image_paths)} 个图像文件")

            all_results = [self.process_image(image_path) for image_path in image_paths]
            self.save_results(all_results, output_path)

            for image_path in image_paths:
                os.remove(image_path)
            os.rmdir(temp_dir)

            return output_path + '.json', output_path + '.txt'
        except Exception as e:
            print(f"处理PDF时发生错误: {str(e)}")
            raise


def main():
    """
    对于同一个两页的pdf识别任务
    仅使用CPU，耗时51.50秒
    使用GPU，耗时11.55秒
    效率提升了347%
    :return:
    """
    processor = PDFOCRProcessor(use_gpu=True, debug=True)

    pdf_path = "./1-2.pdf"
    output_path = "./output/easyocr/result"

    try:
        json_path, txt_path = processor.process_pdf(pdf_path, output_path)
        print(f"处理完成！")
        print(f"JSON结果保存在: {json_path}")
        print(f"文本结果保存在: {txt_path}")
    except Exception as e:
        print(f"处理失败: {str(e)}")

if __name__ == "__main__":
    # check_gpu()
    import time

    start_time = time.time()
    main()
    end_time = time.time()
    execution_time = end_time - start_time
    print(f"函数执行时间: {execution_time:.2f} 秒")
